{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The config attributes {'scaling_factor': 0.18215} were passed to AutoencoderKL, but are not expected and will be ignored. Please verify your config.json configuration file.\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('cuda',\n",
       " PNDMScheduler {\n",
       "   \"_class_name\": \"PNDMScheduler\",\n",
       "   \"_diffusers_version\": \"0.12.1\",\n",
       "   \"beta_end\": 0.012,\n",
       "   \"beta_schedule\": \"scaled_linear\",\n",
       "   \"beta_start\": 0.00085,\n",
       "   \"clip_sample\": false,\n",
       "   \"num_train_timesteps\": 1000,\n",
       "   \"prediction_type\": \"epsilon\",\n",
       "   \"set_alpha_to_one\": false,\n",
       "   \"skip_prk_steps\": true,\n",
       "   \"steps_offset\": 1,\n",
       "   \"trained_betas\": null\n",
       " },\n",
       " CLIPTokenizer(name_or_path='models/diffsion_from_scratch.params/tokenizer', vocab_size=49408, model_max_length=77, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<|startoftext|>', 'eos_token': '<|endoftext|>', 'unk_token': '<|endoftext|>', 'pad_token': '<|endoftext|>'}, clean_up_tokenization_spaces=True),  added_tokens_decoder={\n",
       " \t49406: AddedToken(\"<|startoftext|>\", rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),\n",
       " \t49407: AddedToken(\"<|endoftext|>\", rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),\n",
       " })"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from diffusers import DiffusionPipeline\n",
    "import torch\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "pipeline = DiffusionPipeline.from_pretrained(\n",
    "    'models/diffsion_from_scratch.params', safety_checker=None)\n",
    "\n",
    "scheduler = pipeline.scheduler\n",
    "tokenizer = pipeline.tokenizer\n",
    "\n",
    "del pipeline\n",
    "\n",
    "device, scheduler, tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from datasets import load_dataset\n",
    "\n",
    "# #加载数据集\n",
    "# dataset = load_dataset(path='lansinuote/diffsion_from_scratch', split='train')\n",
    "# dataset.save_to_disk(\"dataset/diffsion_from_scratch\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 833/833 [00:30<00:00, 27.47 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(Dataset({\n",
       "     features: ['pixel_values', 'input_ids'],\n",
       "     num_rows: 833\n",
       " }),\n",
       " {'pixel_values': tensor([[[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           ...,\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       "  \n",
       "          [[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           ...,\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       "  \n",
       "          [[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           ...,\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.]]]),\n",
       "  'input_ids': tensor([49406,   320,  3610,   539,   320,  1901,  9528,   593,   736,  3095,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407])})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_from_disk\n",
    "import torchvision\n",
    "\n",
    "#加载数据集\n",
    "dataset = load_from_disk(\"dataset/diffsion_from_scratch\")\n",
    "\n",
    "#图像增强模块\n",
    "compose = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize(\n",
    "        512, interpolation=torchvision.transforms.InterpolationMode.BILINEAR),\n",
    "    torchvision.transforms.CenterCrop(512),\n",
    "    #torchvision.transforms.RandomHorizontalFlip(),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize([0.5], [0.5]),\n",
    "])\n",
    "\n",
    "\n",
    "def f(data):\n",
    "    #应用图像增强\n",
    "    pixel_values = [compose(i) for i in data['image']]\n",
    "\n",
    "    #文字编码\n",
    "    input_ids = tokenizer.batch_encode_plus(data['text'],\n",
    "                                            padding='max_length',\n",
    "                                            truncation=True,\n",
    "                                            max_length=77).input_ids\n",
    "\n",
    "    return {'pixel_values': pixel_values, 'input_ids': input_ids}\n",
    "\n",
    "\n",
    "dataset = dataset.map(f,\n",
    "                      batched=True,\n",
    "                      batch_size=100,\n",
    "                      num_proc=1,\n",
    "                      remove_columns=['image', 'text'])\n",
    "\n",
    "dataset.set_format(type='torch')\n",
    "\n",
    "dataset, dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(833,\n",
       " {'pixel_values': tensor([[[[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            ...,\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       "  \n",
       "           [[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            ...,\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       "  \n",
       "           [[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            ...,\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.]]]], device='cuda:0'),\n",
       "  'input_ids': tensor([[49406,   320,  1070,  2242,  1312,  7651,  4009,   593,  1205,  3095,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407]], device='cuda:0')})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#定义loader\n",
    "def collate_fn(data):\n",
    "    pixel_values = [i['pixel_values'] for i in data]\n",
    "    input_ids = [i['input_ids'] for i in data]\n",
    "\n",
    "    pixel_values = torch.stack(pixel_values).to(device)\n",
    "    input_ids = torch.stack(input_ids).to(device)\n",
    "\n",
    "    return {'pixel_values': pixel_values, 'input_ids': input_ids}\n",
    "\n",
    "\n",
    "loader = torch.utils.data.DataLoader(dataset,\n",
    "                                     shuffle=True,\n",
    "                                     collate_fn=collate_fn,\n",
    "                                     batch_size=1)\n",
    "\n",
    "len(loader), next(iter(loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The config attributes {'scaling_factor': 0.18215} were passed to AutoencoderKL, but are not expected and will be ignored. Please verify your config.json configuration file.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(AdamW (\n",
       " Parameter Group 0\n",
       "     amsgrad: False\n",
       "     betas: (0.9, 0.999)\n",
       "     capturable: False\n",
       "     differentiable: False\n",
       "     eps: 1e-08\n",
       "     foreach: None\n",
       "     fused: None\n",
       "     lr: 1e-05\n",
       "     maximize: False\n",
       "     weight_decay: 0.01\n",
       " ),\n",
       " MSELoss())"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#加载模型\n",
    "%run encoder.ipynb\n",
    "%run vae.ipynb\n",
    "%run unet.ipynb\n",
    "\n",
    "#准备训练\n",
    "encoder.requires_grad_(False)\n",
    "vae.requires_grad_(False)\n",
    "unet.requires_grad_(True)\n",
    "\n",
    "encoder.eval()\n",
    "vae.eval()\n",
    "unet.train()\n",
    "\n",
    "encoder.to(device)\n",
    "vae.to(device)\n",
    "unet.to(device)\n",
    "\n",
    "optimizer = torch.optim.AdamW(unet.parameters(),\n",
    "                              lr=1e-5,\n",
    "                              betas=(0.9, 0.999),\n",
    "                              weight_decay=0.01,\n",
    "                              eps=1e-8)\n",
    "\n",
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "optimizer, criterion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.0428, device='cuda:0', grad_fn=<MseLossBackward0>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_loss(data):\n",
    "    with torch.no_grad():\n",
    "        #文字编码\n",
    "        #[1, 77] -> [1, 77, 768]\n",
    "        out_encoder = encoder(data['input_ids'])\n",
    "\n",
    "        #抽取图像特征图\n",
    "        #[1, 3, 512, 512] -> [1, 4, 64, 64]\n",
    "        out_vae = vae.encoder(data['pixel_values'])\n",
    "        out_vae = vae.sample(out_vae)\n",
    "\n",
    "        #0.18215 = vae.config.scaling_factor\n",
    "        out_vae = out_vae * 0.18215\n",
    "\n",
    "    #随机数,unet的计算目标\n",
    "    noise = torch.randn_like(out_vae)\n",
    "\n",
    "    #往特征图中添加噪声\n",
    "    #1000 = scheduler.num_train_timesteps\n",
    "    #1 = batch size\n",
    "    noise_step = torch.randint(0, 1000, (1, )).long().to(device)\n",
    "    out_vae_noise = scheduler.add_noise(out_vae, noise, noise_step)\n",
    "\n",
    "    #根据文字信息,把特征图中的噪声计算出来\n",
    "    out_unet = unet(out_vae=out_vae_noise,\n",
    "                    out_encoder=out_encoder,\n",
    "                    time=noise_step)\n",
    "\n",
    "    #计算mse loss\n",
    "    #[1, 4, 64, 64],[1, 4, 64, 64]\n",
    "    return criterion(out_unet, noise)\n",
    "\n",
    "\n",
    "get_loss({\n",
    "    'input_ids': torch.ones(1, 77, device=device).long(),\n",
    "    'pixel_values': torch.randn(1, 3, 512, 512, device=device)\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 11.543325656559318\n"
     ]
    }
   ],
   "source": [
    "def train():\n",
    "    loss_sum = 0\n",
    "    for epoch in range(400):\n",
    "        for i, data in enumerate(loader):\n",
    "            loss = get_loss(data) / 4\n",
    "            loss.backward()\n",
    "            loss_sum += loss.item()\n",
    "\n",
    "            if (epoch * len(loader) + i) % 4 == 0:\n",
    "                torch.nn.utils.clip_grad_norm_(unet.parameters(), 1.0)\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "        if epoch % 10 == 0:\n",
    "            print(epoch, loss_sum)\n",
    "            loss_sum = 0\n",
    "\n",
    "    #torch.save(unet.to('cpu'), 'saves/unet.model')\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import PreTrainedModel, PretrainedConfig\n",
    "\n",
    "\n",
    "#包装类\n",
    "class Model(PreTrainedModel):\n",
    "    config_class = PretrainedConfig\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__(config)\n",
    "        self.unet = unet.to('cpu')\n",
    "\n",
    "# #保存到hub\n",
    "# Model(PretrainedConfig()).push_to_hub(\n",
    "#     repo_id='lansinuote/diffsion_from_scratch.unet',\n",
    "#     use_auth_token=open('/root/hub_token.txt').read().strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "simple",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
